import torch
import torch.nn as nn
import torch.nn.functional as F

class BasicBlock(nn.Module):
    expansion=1
    def __init__(self,in_planes,out_planes,stride=1):
        super(BasicBlock, self).__init__()
        self.conv1=nn.Conv2d(in_planes,out_planes,kernel_size=3,stride=stride,padding=1,bias=False)
        self.bn1=nn.BatchNorm2d(out_planes)
        self.conv2=nn.Conv2d(out_planes,out_planes,kernel_size=3,stride=stride,padding=1,bias=False)
        self.bn2=nn.BatchNorm2d(out_planes)

        self.shortcut=nn.Sequential()

        if stride!=1 or in_planes!=self.expansion*out_planes:
            self.shortcut=nn.Sequential(
                nn.Conv2d(in_planes,self.expansion*out_planes,kernel_size=1,stride=stride,bias=False),
                nn.BatchNorm2d(self.expansion*out_planes)
            )

    def forward(self,x):
        out=F.relu(self.bn1(self.conv1(x)))
        out=self.bn2(self.conv2(out))
        out+=self.shortcut(x)
        out=F.relu(out)
        return out


class Bottleneck(nn.Module):
    expansion=4
    def __init__(self,in_planes,out_planes,stride=1):
        super(Bottleneck, self).__init__()
        self.conv1=nn.Conv2d(in_planes,out_planes,kernel_size=1,bias=False)
        self.bn1=nn.BatchNorm2d(out_planes)
        self.conv2=nn.Conv2d(out_planes,out_planes,kernel_size=3,stride=stride,padding=1,bias=False)
        self.bn2=nn.BatchNorm2d(out_planes)
        self.conv3=nn.Conv2d(out_planes,self.expansion*out_planes,kernel_size=1,bias=False)
        self.bn3=nn.BatchNorm2d(self.expansion*out_planes)

        self.shortcut=nn.Sequential()
        if stride!=1 or in_planes!=self.expansion*out_planes:
            self.shortcut=nn.Sequential(
                nn.Conv2d(in_planes,self.expansion*out_planes,kernel_size=1,stride=stride,bias=False),
                nn.BatchNorm2d(self.expansion*out_planes)
            )
    def forward(self,x):
        out=F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


